package prog8.optimizer

import prog8.ast.IFunctionCall
import prog8.ast.Program
import prog8.ast.base.*
import prog8.ast.expressions.*
import prog8.ast.processing.IAstModifyingVisitor
import prog8.ast.statements.*
import prog8.compiler.target.CompilationTarget
import prog8.functions.BuiltinFunctions


// TODO implement using AstWalker instead of IAstModifyingVisitor
internal class ConstantFoldingOptimizer(private val program: Program, private val errors: ErrorReporter) : IAstModifyingVisitor {
    var optimizationsDone: Int = 0

    override fun visit(decl: VarDecl): Statement {
        // the initializer value can't refer to the variable itself (recursive definition)
        // TODO: use call graph for this?
        if(decl.value?.referencesIdentifiers(decl.name) == true || decl.arraysize?.index?.referencesIdentifiers(decl.name) == true) {
            errors.err("recursive var declaration", decl.position)
            return decl
        }

        if(decl.type==VarDeclType.CONST || decl.type==VarDeclType.VAR) {
            if(decl.isArray){
                if(decl.arraysize==null) {
                    // for arrays that have no size specifier (or a non-constant one) attempt to deduce the size
                    val arrayval = decl.value as? ArrayLiteralValue
                    if(arrayval!=null) {
                        decl.arraysize = ArrayIndex(NumericLiteralValue.optimalInteger(arrayval.value.size, decl.position), decl.position)
                        optimizationsDone++
                    }
                }
                else if(decl.arraysize?.size()==null) {
                    val size = decl.arraysize!!.index.accept(this)
                    if(size is NumericLiteralValue) {
                        decl.arraysize = ArrayIndex(size, decl.position)
                        optimizationsDone++
                    }
                }
            }

            when(decl.datatype) {
                DataType.FLOAT -> {
                    // vardecl: for scalar float vars, promote constant integer initialization values to floats
                    val litval = decl.value as? NumericLiteralValue
                    if (litval!=null && litval.type in IntegerDatatypes) {
                        val newValue = NumericLiteralValue(DataType.FLOAT, litval.number.toDouble(), litval.position)
                        decl.value = newValue
                        optimizationsDone++
                        return super.visit(decl)
                    }
                }
                DataType.ARRAY_UB, DataType.ARRAY_B, DataType.ARRAY_UW, DataType.ARRAY_W -> {
                    val numericLv = decl.value as? NumericLiteralValue
                    val rangeExpr = decl.value as? RangeExpr
                    if(rangeExpr!=null) {
                        // convert the initializer range expression to an actual array
                        val declArraySize = decl.arraysize?.size()
                        if(declArraySize!=null && declArraySize!=rangeExpr.size())
                            errors.err("range expression size doesn't match declared array size", decl.value?.position!!)
                        val constRange = rangeExpr.toConstantIntegerRange()
                        if(constRange!=null) {
                            val eltType = rangeExpr.inferType(program).typeOrElse(DataType.UBYTE)
                            if(eltType in ByteDatatypes) {
                                decl.value = ArrayLiteralValue(InferredTypes.InferredType.known(decl.datatype),
                                        constRange.map { NumericLiteralValue(eltType, it.toShort(), decl.value!!.position) }.toTypedArray(),
                                        position = decl.value!!.position)
                            } else {
                                decl.value = ArrayLiteralValue(InferredTypes.InferredType.known(decl.datatype),
                                        constRange.map { NumericLiteralValue(eltType, it, decl.value!!.position) }.toTypedArray(),
                                        position = decl.value!!.position)
                            }
                            decl.value!!.linkParents(decl)
                            optimizationsDone++
                            return super.visit(decl)
                        }
                    }
                    if(numericLv!=null && numericLv.type== DataType.FLOAT)
                        errors.err("arraysize requires only integers here", numericLv.position)
                    val size = decl.arraysize?.size() ?: return decl
                    if (rangeExpr==null && numericLv!=null) {
                        // arraysize initializer is empty or a single int, and we know the size; create the arraysize.
                        val fillvalue = numericLv.number.toInt()
                        when(decl.datatype){
                            DataType.ARRAY_UB -> {
                                if(fillvalue !in 0..255)
                                    errors.err("ubyte value overflow", numericLv.position)
                            }
                            DataType.ARRAY_B -> {
                                if(fillvalue !in -128..127)
                                    errors.err("byte value overflow", numericLv.position)
                            }
                            DataType.ARRAY_UW -> {
                                if(fillvalue !in 0..65535)
                                    errors.err("uword value overflow", numericLv.position)
                            }
                            DataType.ARRAY_W -> {
                                if(fillvalue !in -32768..32767)
                                    errors.err("word value overflow", numericLv.position)
                            }
                            else -> {}
                        }
                        // create the array itself, filled with the fillvalue.
                        val array = Array(size) {fillvalue}.map { NumericLiteralValue(ArrayElementTypes.getValue(decl.datatype), it, numericLv.position) as Expression}.toTypedArray()
                        val refValue = ArrayLiteralValue(InferredTypes.InferredType.known(decl.datatype), array, position = numericLv.position)
                        decl.value = refValue
                        refValue.parent=decl
                        optimizationsDone++
                        return super.visit(decl)
                    }
                }
                DataType.ARRAY_F  -> {
                    val size = decl.arraysize?.size() ?: return decl
                    val litval = decl.value as? NumericLiteralValue
                    if(litval==null) {
                        // there's no initialization value, but the size is known, so we're ok.
                        return super.visit(decl)
                    } else {
                        // arraysize initializer is a single int, and we know the size.
                        val fillvalue = litval.number.toDouble()
                        if (fillvalue < CompilationTarget.machine.FLOAT_MAX_NEGATIVE || fillvalue > CompilationTarget.machine.FLOAT_MAX_POSITIVE)
                            errors.err("float value overflow", litval.position)
                        else {
                            // create the array itself, filled with the fillvalue.
                            val array = Array(size) {fillvalue}.map { NumericLiteralValue(DataType.FLOAT, it, litval.position) as Expression}.toTypedArray()
                            val refValue = ArrayLiteralValue(InferredTypes.InferredType.known(DataType.ARRAY_F), array, position = litval.position)
                            decl.value = refValue
                            refValue.parent=decl
                            optimizationsDone++
                            return super.visit(decl)
                        }
                    }
                }
                else -> {
                    // nothing to do for this type
                    // this includes strings and structs
                }
            }
        }

        val declValue = decl.value
        if(declValue!=null && decl.type==VarDeclType.VAR
                && declValue is NumericLiteralValue && !declValue.inferType(program).istype(decl.datatype)) {
            // cast the numeric literal to the appropriate datatype of the variable
            decl.value = declValue.cast(decl.datatype)
        }

        return super.visit(decl)
    }

    /**
     * replace identifiers that refer to const value, with the value itself (if it's a simple type)
     */
    override fun visit(identifier: IdentifierReference): Expression {
        // don't replace when it's an assignment target or loop variable
        if(identifier.parent is AssignTarget)
            return identifier
        var forloop = identifier.parent as? ForLoop
        if(forloop==null)
            forloop = identifier.parent.parent as? ForLoop
        if(forloop!=null && identifier===forloop.loopVar)
            return identifier

        val cval = identifier.constValue(program) ?: return identifier
        return when (cval.type) {
            in NumericDatatypes -> {
                val copy = NumericLiteralValue(cval.type, cval.number, identifier.position)
                copy.parent = identifier.parent
                copy
            }
            in PassByReferenceDatatypes -> throw FatalAstException("pass-by-reference type should not be considered a constant")
            else -> identifier
        }
    }

    override fun visit(functionCall: FunctionCall): Expression {
        super.visit(functionCall)
        typeCastConstArguments(functionCall)
        return functionCall.constValue(program) ?: functionCall
    }

    override fun visit(functionCallStatement: FunctionCallStatement): Statement {
        super.visit(functionCallStatement)
        typeCastConstArguments(functionCallStatement)
        return functionCallStatement
    }

    private fun typeCastConstArguments(functionCall: IFunctionCall) {
        if(functionCall.target.nameInSource.size==1) {
            val builtinFunction = BuiltinFunctions[functionCall.target.nameInSource.single()]
            if(builtinFunction!=null) {
                // match the arguments of a builtin function signature.
                for(arg in functionCall.args.withIndex().zip(builtinFunction.parameters)) {
                    val possibleDts = arg.second.possibleDatatypes
                    val argConst = arg.first.value.constValue(program)
                    if(argConst!=null && argConst.type !in possibleDts) {
                        val convertedValue = argConst.cast(possibleDts.first())
                        functionCall.args[arg.first.index] = convertedValue
                        optimizationsDone++
                    }
                }
                return
            }
        }
        // match the arguments of a subroutine.
        val subroutine = functionCall.target.targetSubroutine(program.namespace)
        if(subroutine!=null) {
            // if types differ, try to typecast constant arguments to the function call to the desired data type of the parameter
            for(arg in functionCall.args.withIndex().zip(subroutine.parameters)) {
                val expectedDt = arg.second.type
                val argConst = arg.first.value.constValue(program)
                if(argConst!=null && argConst.type!=expectedDt) {
                    val convertedValue = argConst.cast(expectedDt)
                    functionCall.args[arg.first.index] = convertedValue
                    optimizationsDone++
                }
            }
        }
    }

    override fun visit(memread: DirectMemoryRead): Expression {
        // @( &thing )  -->  thing
        val addrOf = memread.addressExpression as? AddressOf
        if(addrOf!=null)
            return super.visit(addrOf.identifier)
        return super.visit(memread)
    }

    /**
     * Try to accept a unary prefix expression.
     * Compile-time constant sub expressions will be evaluated on the spot.
     * For instance, the expression for "- 4.5" will be optimized into the float literal -4.5
     */
    override fun visit(expr: PrefixExpression): Expression {
        val prefixExpr=super.visit(expr)
        if(prefixExpr !is PrefixExpression)
            return prefixExpr

        val subexpr = prefixExpr.expression
        if (subexpr is NumericLiteralValue) {
            // accept prefixed literal values (such as -3, not true)
            return when (prefixExpr.operator) {
                "+" -> subexpr
                "-" -> when (subexpr.type) {
                    in IntegerDatatypes -> {
                        optimizationsDone++
                        NumericLiteralValue.optimalNumeric(-subexpr.number.toInt(), subexpr.position)
                    }
                    DataType.FLOAT -> {
                        optimizationsDone++
                        NumericLiteralValue(DataType.FLOAT, -subexpr.number.toDouble(), subexpr.position)
                    }
                    else -> throw ExpressionError("can only take negative of int or float", subexpr.position)
                }
                "~" -> when (subexpr.type) {
                    in IntegerDatatypes -> {
                        optimizationsDone++
                        NumericLiteralValue.optimalNumeric(subexpr.number.toInt().inv(), subexpr.position)
                    }
                    else -> throw ExpressionError("can only take bitwise inversion of int", subexpr.position)
                }
                "not" -> {
                    optimizationsDone++
                    NumericLiteralValue.fromBoolean(subexpr.number.toDouble() == 0.0, subexpr.position)
                }
                else -> throw ExpressionError(prefixExpr.operator, subexpr.position)
            }
        }
        return prefixExpr
    }

    /**
     * Try to accept a binary expression.
     * Compile-time constant sub expressions will be evaluated on the spot.
     * For instance, "9 * (4 + 2)" will be optimized into the integer literal 54.
     *
     * More complex stuff: reordering to group constants:
     * If one of our operands is a Constant,
     *   and the other operand is a Binary expression,
     *   and one of ITS operands is a Constant,
     *   and ITS other operand is NOT a Constant,
     *   ...it may be possible to rewrite the expression to group the two Constants together,
     *      to allow them to be const-folded away.
     *
     *  examples include:
     *        (X / c1) * c2  ->  X / (c2/c1)
     *        (X + c1) - c2  ->  X + (c1-c2)
     */
    override fun visit(expr: BinaryExpression): Expression {
        super.visit(expr)

        if(expr.left is StringLiteralValue || expr.left is ArrayLiteralValue
                || expr.right is StringLiteralValue || expr.right is ArrayLiteralValue)
            throw FatalAstException("binexpr with reference litval instead of numeric")

        val leftconst = expr.left.constValue(program)
        val rightconst = expr.right.constValue(program)

        val subExpr: BinaryExpression? = when {
            leftconst!=null -> expr.right as? BinaryExpression
            rightconst!=null -> expr.left as? BinaryExpression
            else -> null
        }
        if(subExpr!=null) {
            val subleftconst = subExpr.left.constValue(program)
            val subrightconst = subExpr.right.constValue(program)
            if ((subleftconst != null && subrightconst == null) || (subleftconst==null && subrightconst!=null)) {
                // try reordering.
                return groupTwoConstsTogether(expr, subExpr,
                        leftconst != null, rightconst != null,
                        subleftconst != null, subrightconst != null)
            }
        }

        // const fold when both operands are a const
        return when {
            leftconst != null && rightconst != null -> {
                optimizationsDone++
                val evaluator = ConstExprEvaluator()
                evaluator.evaluate(leftconst, expr.operator, rightconst)
            }

            else -> expr
        }
    }

    private fun groupTwoConstsTogether(expr: BinaryExpression,
                                       subExpr: BinaryExpression,
                                       leftIsConst: Boolean,
                                       rightIsConst: Boolean,
                                       subleftIsConst: Boolean,
                                       subrightIsConst: Boolean): Expression
    {
        // todo: this implements only a small set of possible reorderings at this time
        if(expr.operator==subExpr.operator) {
            // both operators are the isSameAs.
            // If + or *,  we can simply swap the const of expr and Var in subexpr.
            if(expr.operator=="+" || expr.operator=="*") {
                if(leftIsConst) {
                    if(subleftIsConst)
                        expr.left = subExpr.right.also { subExpr.right = expr.left }
                    else
                        expr.left = subExpr.left.also { subExpr.left = expr.left }
                } else {
                    if(subleftIsConst)
                        expr.right = subExpr.right.also {subExpr.right = expr.right }
                    else
                        expr.right = subExpr.left.also { subExpr.left = expr.right }
                }
                optimizationsDone++
                return expr
            }

            // If - or /,  we simetimes must reorder more, and flip operators (- -> +, / -> *)
            if(expr.operator=="-" || expr.operator=="/") {
                optimizationsDone++
                if(leftIsConst) {
                    return if(subleftIsConst) {
                        val tmp = subExpr.right
                        subExpr.right = subExpr.left
                        subExpr.left = expr.left
                        expr.left = tmp
                        expr.operator = if(expr.operator=="-") "+" else "*"
                        expr
                    } else
                        BinaryExpression(
                                BinaryExpression(expr.left, if (expr.operator == "-") "+" else "*", subExpr.right, subExpr.position),
                                expr.operator, subExpr.left, expr.position)
                } else {
                    return if(subleftIsConst) {
                        expr.right = subExpr.right.also { subExpr.right = expr.right }
                        expr
                    } else
                        BinaryExpression(
                                subExpr.left, expr.operator,
                                BinaryExpression(expr.right, if (expr.operator == "-") "+" else "*", subExpr.right, subExpr.position),
                                expr.position)
                }
            }
            return expr

        }
        else
        {

            if(expr.operator=="/" && subExpr.operator=="*") {
                optimizationsDone++
                if(leftIsConst) {
                    return if(subleftIsConst) {
                        // C1/(C2*V) -> (C1/C2)/V
                        BinaryExpression(
                                BinaryExpression(expr.left, "/", subExpr.left, subExpr.position),
                                "/",
                                subExpr.right, expr.position)
                    } else {
                        // C1/(V*C2) -> (C1/C2)/V
                        BinaryExpression(
                                BinaryExpression(expr.left, "/", subExpr.right, subExpr.position),
                                "/",
                                subExpr.left, expr.position)
                    }
                } else {
                    return if(subleftIsConst) {
                        // (C1*V)/C2 -> (C1/C2)*V
                        BinaryExpression(
                                BinaryExpression(subExpr.left, "/", expr.right, subExpr.position),
                                "*",
                                subExpr.right, expr.position)
                    } else {
                        // (V*C1)/C2 -> (C1/C2)*V
                        BinaryExpression(
                                BinaryExpression(subExpr.right, "/", expr.right, subExpr.position),
                                "*",
                                subExpr.left, expr.position)
                    }
                }
            }
            else if(expr.operator=="*" && subExpr.operator=="/") {
                optimizationsDone++
                if(leftIsConst) {
                    return if(subleftIsConst) {
                        // C1*(C2/V) -> (C1*C2)/V
                        BinaryExpression(
                                BinaryExpression(expr.left, "*", subExpr.left, subExpr.position),
                                "/",
                                subExpr.right, expr.position)
                    } else {
                        // C1*(V/C2) -> (C1/C2)*V
                        BinaryExpression(
                                BinaryExpression(expr.left, "/", subExpr.right, subExpr.position),
                                "*",
                                subExpr.left, expr.position)
                    }
                } else {
                    return if(subleftIsConst) {
                        // (C1/V)*C2 -> (C1*C2)/V
                        BinaryExpression(
                                BinaryExpression(subExpr.left, "*", expr.right, subExpr.position),
                                "/",
                                subExpr.right, expr.position)
                    } else {
                        // (V/C1)*C2 -> (C1/C2)*V
                        BinaryExpression(
                                BinaryExpression(expr.right, "/", subExpr.right, subExpr.position),
                                "*",
                                subExpr.left, expr.position)
                    }
                }
            }
            else if(expr.operator=="+" && subExpr.operator=="-") {
                optimizationsDone++
                if(leftIsConst){
                    return if(subleftIsConst){
                        // c1+(c2-v)  ->  (c1+c2)-v
                        BinaryExpression(
                                BinaryExpression(expr.left, "+", subExpr.left, subExpr.position),
                                "-",
                                subExpr.right, expr.position)
                    } else {
                        // c1+(v-c2)  ->  v+(c1-c2)
                        BinaryExpression(
                                BinaryExpression(expr.left, "-", subExpr.right, subExpr.position),
                                "+",
                                subExpr.left, expr.position)
                    }
                } else {
                    return if(subleftIsConst) {
                        // (c1-v)+c2  ->  (c1+c2)-v
                        BinaryExpression(
                                BinaryExpression(subExpr.left, "+", expr.right, subExpr.position),
                                "-",
                                subExpr.right, expr.position)
                    } else {
                        // (v-c1)+c2  ->  v+(c2-c1)
                        BinaryExpression(
                                BinaryExpression(expr.right, "-", subExpr.right, subExpr.position),
                                "+",
                                subExpr.left, expr.position)
                    }
                }
            }
            else if(expr.operator=="-" && subExpr.operator=="+") {
                optimizationsDone++
                if(leftIsConst) {
                    return if(subleftIsConst) {
                        // c1-(c2+v)  ->  (c1-c2)-v
                        BinaryExpression(
                                BinaryExpression(expr.left, "-", subExpr.left, subExpr.position),
                                "-",
                                subExpr.right, expr.position)
                    } else {
                        // c1-(v+c2)  ->  (c1-c2)-v
                        BinaryExpression(
                                BinaryExpression(expr.left, "-", subExpr.right, subExpr.position),
                                "-",
                                subExpr.left, expr.position)
                    }
                } else {
                    return if(subleftIsConst) {
                        // (c1+v)-c2  ->  v+(c1-c2)
                        BinaryExpression(
                                BinaryExpression(subExpr.left, "-", expr.right, subExpr.position),
                                "+",
                                subExpr.right, expr.position)
                    } else {
                        // (v+c1)-c2  ->  v+(c1-c2)
                        BinaryExpression(
                                BinaryExpression(subExpr.right, "-", expr.right, subExpr.position),
                                "+",
                                subExpr.left, expr.position)
                    }
                }
            }

            return expr
        }
    }

    override fun visit(forLoop: ForLoop): Statement {

        fun adjustRangeDt(rangeFrom: NumericLiteralValue, targetDt: DataType, rangeTo: NumericLiteralValue, stepLiteral: NumericLiteralValue?, range: RangeExpr): RangeExpr {
            val newFrom: NumericLiteralValue
            val newTo: NumericLiteralValue
            try {
                newFrom = rangeFrom.cast(targetDt)
                newTo = rangeTo.cast(targetDt)
            } catch (x: ExpressionError) {
                return range
            }
            val newStep: Expression = try {
                stepLiteral?.cast(targetDt)?: range.step
            } catch(ee: ExpressionError) {
                range.step
            }
            return RangeExpr(newFrom, newTo, newStep, range.position)
        }

        val forLoop2 = super.visit(forLoop) as ForLoop

        // check if we need to adjust an array literal to the loop variable's datatype
        val array = forLoop2.iterable as? ArrayLiteralValue
        if(array!=null) {
            val loopvarDt: DataType = when {
                forLoop.loopVar!=null -> forLoop.loopVar!!.inferType(program).typeOrElse(DataType.UBYTE)
                forLoop.loopRegister!=null -> DataType.UBYTE
                else -> throw FatalAstException("weird for loop")
            }

            val arrayType = when(loopvarDt) {
                DataType.UBYTE -> DataType.ARRAY_UB
                DataType.BYTE -> DataType.ARRAY_B
                DataType.UWORD -> DataType.ARRAY_UW
                DataType.WORD -> DataType.ARRAY_W
                DataType.FLOAT -> DataType.ARRAY_F
                else -> throw FatalAstException("invalid array elt type")
            }
            val array2 = array.cast(arrayType)
            if(array2!=null && array2!==array) {
                forLoop2.iterable = array2
                array2.linkParents(forLoop2)
            }
        }

        // adjust the datatype of a range expression in for loops to the loop variable.
        val iterableRange = forLoop2.iterable as? RangeExpr ?: return forLoop2
        val rangeFrom = iterableRange.from as? NumericLiteralValue
        val rangeTo = iterableRange.to as? NumericLiteralValue
        if(rangeFrom==null || rangeTo==null) return forLoop2

        val loopvar = forLoop2.loopVar?.targetVarDecl(program.namespace)
        if(loopvar!=null) {
            val stepLiteral = iterableRange.step as? NumericLiteralValue
            when(loopvar.datatype) {
                DataType.UBYTE -> {
                    if(rangeFrom.type!= DataType.UBYTE) {
                        // attempt to translate the iterable into ubyte values
                        forLoop2.iterable = adjustRangeDt(rangeFrom, loopvar.datatype, rangeTo, stepLiteral, iterableRange)
                    }
                }
                DataType.BYTE -> {
                    if(rangeFrom.type!= DataType.BYTE) {
                        // attempt to translate the iterable into byte values
                        forLoop2.iterable = adjustRangeDt(rangeFrom, loopvar.datatype, rangeTo, stepLiteral, iterableRange)
                    }
                }
                DataType.UWORD -> {
                    if(rangeFrom.type!= DataType.UWORD) {
                        // attempt to translate the iterable into uword values
                        forLoop2.iterable = adjustRangeDt(rangeFrom, loopvar.datatype, rangeTo, stepLiteral, iterableRange)
                    }
                }
                DataType.WORD -> {
                    if(rangeFrom.type!= DataType.WORD) {
                        // attempt to translate the iterable into word values
                        forLoop2.iterable = adjustRangeDt(rangeFrom, loopvar.datatype, rangeTo, stepLiteral, iterableRange)
                    }
                }
                else -> throw FatalAstException("invalid loopvar datatype $loopvar")
            }
        }
        return forLoop2
    }

    override fun visit(arrayLiteral: ArrayLiteralValue): Expression {
        // because constant folding can result in arrays that are now suddenly capable
        // of telling the type of all their elements (for instance, when they contained -2 which
        // was a prefix expression earlier), we recalculate the array's datatype.
        val array = super.visit(arrayLiteral)
        if(array is ArrayLiteralValue) {
            if(array.type.isKnown)
                return array
            val arrayDt = array.guessDatatype(program)
            if(arrayDt.isKnown) {
                val newArray = arrayLiteral.cast(arrayDt.typeOrElse(DataType.STRUCT))
                if(newArray!=null)
                    return newArray
            }
        }
        return array
    }
}
